# 优化后的 sftp_util.py
import os
import paramiko
from config.config import remote_server, remote_port, remote_username, remote_password, ppa_prefix_to_api, \
    ppa_shared_folder_path, ppa_prefix_to_folder
import logging

from util.initial_log import initial_log_config


class SFTPUtil:
    """
    SFTP工具类，用于管理SFTP连接和文件传输操作
    """

    def __init__(self, server=None, port=None, username=None, password=None):
        """
        初始化SFTP工具类

        Args:
            server (str, optional): SFTP服务器地址，默认使用config中的配置
            port (int, optional): SFTP端口，默认使用config中的配置
            username (str, optional): 用户名，默认使用config中的配置
            password (str, optional): 密码，默认使用config中的配置
        """
        self.server = server or remote_server
        self.port = port or remote_port
        self.username = username or remote_username
        self.password = password or remote_password
        self.sftp = None
        self.transport = None

    def connect(self):
        """
        建立SFTP连接

        Returns:
            bool: 连接是否成功
        """
        try:
            self.transport = paramiko.Transport((str(self.server), int(self.port)))
            self.transport.connect(username=str(self.username), password=str(self.password))
            self.sftp = paramiko.SFTPClient.from_transport(self.transport)
            logging.info(f"SFTP connection established successfully to {self.server}:{self.port}")
            return True
        except Exception as e:
            logging.error(f"Failed to establish SFTP connection to {self.server}:{self.port}: {e}")
            return False

    def disconnect(self):
        """
        关闭SFTP连接
        """
        try:
            if self.sftp:
                self.sftp.close()
            if self.transport:
                self.transport.close()
            logging.info(f"SFTP connection to {self.server}:{self.port} closed successfully")
        except Exception as e:
            logging.error(f"Error closing SFTP connection to {self.server}:{self.port}: {e}")

    def _check_remote_dir(self, path):
        """
        检查远程目录是否存在，如果不存在则报错

        Args:
            path (str): 远程目录路径

        Returns:
            bool: 目录是否存在
        """
        # 检查是否已建立连接
        if not self.sftp:
            logging.error("SFTP connection not established. Please call connect() first.")
            return False

        try:
            # 确保路径格式正确，去除可能的多余字符
            clean_path = path.rstrip('/,')
            self.sftp.stat(clean_path)
            return True
        except FileNotFoundError:
            logging.error(f"Remote directory does not exist: {path}")
            return False

    def upload_file(self, local_path, remote_path):
        """
        通过sftp上传单个文件，检查远程目录是否存在

        Args:
            local_path (str): 本地文件路径
            remote_path (str): 远程文件路径

        Returns:
            bool: 上传是否成功
        """
        if not os.path.exists(local_path):
            logging.warning(f"Local file does not exist: {local_path}")
            return False

        # 检查是否已建立连接
        if not self.sftp:
            logging.error("SFTP connection not established. Please call connect() first.")
            return False

        remote_dir = os.path.dirname(remote_path)

        # 检查远程目录是否存在，不存在则报错
        if not self._check_remote_dir(remote_dir):
            logging.error(f"Cannot upload file. Remote directory does not exist: {remote_dir}")
            return False

        try:
            # 确保远程路径格式正确
            clean_remote_path = remote_path.rstrip('/,')
            self.sftp.put(local_path, clean_remote_path)
            logging.info(f"SFTP uploaded {local_path} -> {clean_remote_path}")
            return True
        except Exception as e:
            logging.error(f"Failed to upload {local_path}: {e}")
            return False

    def upload_files_by_prefix(self, source_directory, remote_base_dir, prefix_to_target_folder_map=ppa_prefix_to_folder):
        """
        遍历源目录文件，按前缀上传到远程对应目录

        Args:
            source_directory (str): 本地源目录路径
            prefix_to_target_folder_map (dict): 文件前缀到远程目录的映射
            remote_base_dir (str): 远程基础目录路径

        Returns:
            tuple: (成功上传数量, 失败数量)
        """
        # 检查是否已建立连接
        if not self.sftp:
            logging.error("SFTP connection not established. Please call connect() first.")
            return 0, 0

        if not os.path.exists(source_directory):
            logging.warning(f"Source folder not exist: {source_directory}")
            return 0, 0

        uploaded_count = 0
        failed_count = 0

        for filename in os.listdir(source_directory):
            source_file_path = os.path.join(source_directory, filename)
            if os.path.isfile(source_file_path):
                matched = False
                for prefix, sub_folder in prefix_to_target_folder_map.items():
                    if filename.startswith(prefix):
                        # 确保路径格式正确
                        remote_dir = os.path.join(remote_base_dir, sub_folder).replace('\\', '/').rstrip('/,')
                        remote_file_path = os.path.join(remote_dir, filename).replace('\\', '/').rstrip('/,')

                        # 检查远程目录是否存在
                        if not self._check_remote_dir(remote_dir):
                            logging.error(f"Skipping upload. Remote directory does not exist: {remote_dir}")
                            failed_count += 1
                            matched = True
                            break

                        try:
                            if self.upload_file(source_file_path, remote_file_path):
                                uploaded_count += 1
                            else:
                                failed_count += 1
                            matched = True
                            break
                        except Exception as e:
                            logging.error(f"Upload '{filename}' failed: {e}")
                            failed_count += 1
                if not matched:
                    logging.info(f"File '{filename}' not match any prefix, skip upload.")

        logging.info(f"Upload summary: {uploaded_count} successful, {failed_count} failed")
        return uploaded_count, failed_count

    def upload_calculation_files(self, source_directory, remote_base_dir=None):
        """
        上传计算文件的便捷方法

        Args:
            source_directory (str): 本地源目录路径
            remote_base_dir (str, optional): 远程基础目录路径，默认使用配置中的路径

        Returns:
            bool: 是否成功完成上传
        """
        # 如果未指定远程基础目录，则使用配置中的默认路径
        if remote_base_dir is None:
            remote_base_dir = f"{ppa_shared_folder_path}/in/data"

        # 确保基础路径格式正确
        remote_base_dir = remote_base_dir.rstrip('/,')

        # 建立连接
        if not self.connect():
            logging.error("Failed to establish SFTP connection")
            return False

        try:
            # 执行上传
            uploaded_count, failed_count = self.upload_files_by_prefix(
                source_directory,
                remote_base_dir
            )

            # 记录上传结果
            if failed_count == 0:
                logging.info(f"All files uploaded successfully. Total: {uploaded_count}")
                return True
            else:
                logging.warning(f"Upload completed with errors. Success: {uploaded_count}, Failed: {failed_count}")
                return failed_count == 0
        finally:
            # 确保关闭连接
            self.disconnect()

    def delete_ack_files(self, remote_directory=None):
        """
        删除指定目录下所有的后缀为.ack的文件，包括子目录中的.ack文件

        Args:
            remote_directory (str): 远程目录路径

        Returns:
            tuple: (成功删除数量, 失败数量)
        """

        if remote_directory is None:
            remote_directory = f"{ppa_shared_folder_path}/in/data"

        # 建立连接
        if not self.connect():
            logging.error("Failed to establish SFTP connection")
            return 0, 0

        try:
            # 检查远程目录是否存在
            if not self._check_remote_dir(remote_directory):
                logging.error(f"Cannot delete .ack files. Remote directory does not exist: {remote_directory}")
                return 0, 0

            # 递归遍历目录查找并删除.ack文件
            deleted_count, failed_count = self._delete_ack_files_recursive(remote_directory)
            logging.info(f"Delete .ack files summary: {deleted_count} successful, {failed_count} failed")
            return deleted_count, failed_count
        except Exception as e:
            logging.error(f"Error during deleting .ack files in {remote_directory}: {e}")
            return 0, 0
        finally:
            # 确保关闭连接
            self.disconnect()

    def _delete_ack_files_recursive(self, remote_directory):
        """
        递归删除目录中的.ack文件

        Args:
            remote_directory (str): 远程目录路径

        Returns:
            tuple: (成功删除数量, 失败数量)
        """
        deleted_count = 0
        failed_count = 0

        try:
            # 列出目录中的所有文件和文件夹
            items = self.sftp.listdir_attr(remote_directory)

            for item in items:
                item_path = f"{remote_directory}/{item.filename}".replace('\\', '/').replace('//', '/')

                # 如果是目录，递归处理
                if item.st_mode & 0o040000:  # 检查是否为目录
                    sub_deleted, sub_failed = self._delete_ack_files_recursive(item_path)
                    deleted_count += sub_deleted
                    failed_count += sub_failed
                # 如果是.ack文件，删除它
                elif item.filename.endswith('.ack'):
                    try:
                        self.sftp.remove(item_path)
                        logging.info(f"Deleted .ack file: {item_path}")
                        deleted_count += 1
                    except Exception as e:
                        logging.error(f"Failed to delete .ack file {item_path}: {e}")
                        failed_count += 1

            return deleted_count, failed_count
        except Exception as e:
            logging.error(f"Error listing directory {remote_directory}: {e}")
            raise

    def delete_file(self, remote_file_path):
        """
        删除指定的远程文件

        Args:
            remote_file_path (str): 远程文件的绝对路径

        Returns:
            bool: 删除是否成功
        """
        # 建立连接
        if not self.connect():
            logging.error("Failed to establish SFTP connection")
            return False

        try:
            # 检查文件是否存在
            try:
                self.sftp.stat(remote_file_path)
            except FileNotFoundError:
                logging.warning(f"Remote file does not exist: {remote_file_path}")
                return False

            # 删除文件
            self.sftp.remove(remote_file_path)
            logging.info(f"Successfully deleted remote file: {remote_file_path}")
            return True
        except Exception as e:
            logging.error(f"Failed to delete remote file {remote_file_path}: {e}")
            return False
        finally:
            # 确保关闭连接
            self.disconnect()

if __name__ == '__main__':
    initial_log_config()
    source_directory = r'F:\WorkSpace\pycharam_workspace\ppa_auto\data\calculation\happy_case\low_case\20241231'

    # 使用便捷方法上传文件
    sftp_util = SFTPUtil()
    sftp_util.delete_ack_files()

